package graph;

import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * @BelongsProject: leetcode
 * @BelongsPackage: graph
 * @author: 肖
 * @date: 2022/1/30 10:40
 * @Description: Prim算法 从点开始
 * @since JDK 11
 */

public class Prim {
    public static Set<Edge> primMST(Graph graph) {
//        小根堆
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(
                (o1, o2) -> o1.weight - o2.weight
        );
        HashSet<NodeG> set = new HashSet<>();
        Set<Edge> result = new HashSet<>();
        for (NodeG node : graph.nodes.values()) {
            if (!set.contains(node)) {
                set.add(node);
                for (Edge edge : node.edges) {
                    priorityQueue.add(edge);
                }
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll();
                    NodeG toNode = edge.to;
                    if (!set.contains(toNode)) {
                        set.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge : toNode.edges) {
                            priorityQueue.add(nextEdge);
                        }
                    }
                }
            }

        }
        return result;
    }
}
